# coding=utf-8
# @author:      ChengJing
# @name:        ImproveFeatures.py
# @datetime:    2021/12/10 17:45
# @software:    PyCharm
# @description: 在时序输入数据的基础上，以注意力机制的方式添加流量数据，增加模型定位的能力(注意力数据添加到开始和结束时)

import torch
import torch.nn as nn


class LP(nn.Module):
    """
    添加监测点之间的位置信息
    """
    def __init__(self, features, sequence, sensors, adj):
        """
        Args:
            features: 权重对应的特征数
            sequence: 时间序列的长度
            sensors: 压力监测点的数量
            adj: 压力监测点之间的邻接矩阵
        """
        super(LP, self).__init__()
        self.weight = nn.Parameter(torch.FloatTensor(features*sequence, sensors),requires_grad=True)
        torch.nn.init.xavier_uniform_(self.weight)
        self.adj = adj
        self.f = nn.Linear(sensors, features)

    def forward(self, x):
        """
        前向传递函数
        """
        nx = x.reshape((x.shape[0], -1))
        out = nx.mm(self.weight)
        out = self.f(out.mm(self.adj))
        return out.unsqueeze(dim=1)*x


class AQ(nn.Module):
    """
    以注意力机制的方式添加流量数据，增强数据
    """
    def __init__(self, qsensors, hidden, output):
        """
        Args:
            qsensors: 流量监测点的数量
            hidden: 隐含层节点数
            output: 输出特征数
        """
        super(AQ, self).__init__()
        self.f = nn.Sequential(
            nn.Linear(qsensors, hidden),
            nn.ReLU(),
            nn.Linear(hidden, output),
            nn.Sigmoid()
        )

    def forward(self, x, q):
        """
        前向传递函数
        """
        w = self.f(q)
        return x * w.unsqueeze(dim=1)


class LAQ(nn.Module):
    """
    添加监测点位置信息的流量数据注意力机制
    原理：
        1. 输入数据：G：流量监测点和压力监测点之间的连接关系，shape：N*M； Q：流量监测点的流量数据，shape：1*N；N表示流量监测点的数量，M表示压力监测点的数量。
        第一步：Q*G表示有选择性的学习流量特征
        第二步：添加可学习参数，深入学习流量数据
    """
    def __init__(self, g, sensors, hidden):
        """
        Args:
            G: 流量监测点和压力监测点之间的连接关系，shape：N*M
            sensors: 压力监测点的数量
            hidden: 隐含层节点数
        """
        super(LAQ, self).__init__()
        self.f = nn.Sequential(
            nn.Linear(sensors, hidden),
            nn.ReLU(),
            nn.Linear(hidden, sensors),
            nn.Sigmoid()
        )
        self.g = g

    def forward(self, x, q):
        """
        前向传递函数
        """
        w = q.mm(self.g)
        w = self.f(w)
        return x * w.unsqueeze(dim=1)


if __name__ == '__main__':
    x = torch.rand((10, 20, 6))
    q = torch.rand((10, 3))
    g = torch.rand((3, 6))

    lp = LAQ(g, 6, 36)
    lp(x, q)
    y = x.sum()
    y.requires_grad_(True)
    y.backward()
